
from typing import List
import os
import datetime
from PIL import Image,ImageDraw
import numpy as np
# import folder_paths
import tensorflow as tf
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from modelscope.outputs import OutputKeys
from alioffline_utils_img import *
from alioffline_utils_mask import png_to_mask

# comfy_path = os.path.dirname(folder_paths.__file__)
current_directory = os.path.dirname(__file__)
comfy_path =  os.path.abspath(os.path.join(current_directory, '..', '..'))
models_path = os.path.join(comfy_path, "models","modelscope")
# model_path = os.path.join(models_path, "damo","cv_resnet50_face-detection_retinaface")

class AliOffline_Face_Detection:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ('IMAGE', {}),
            }
        }
    CATEGORY = "CXH"
    FUNCTION = "sample"
    RETURN_NAMES = ("mask",)
    RETURN_TYPES = ("IMAGE",)

    def sample(self,image):
        os.environ['MODELSCOPE_CACHE']=models_path
        res_images = []
        model = pipeline(
            Tasks.face_detection, model="damo/cv_resnet50_face-detection_retinaface"
        )
        for item in image:
          image_item = tensor2pil(item)
          result = model(image_item)
          print(result)
          # {'scores': [0.9985374212265015], 'boxes': [[135.4436798095703, 26.41049575805664, 188.53573608398438, 93.52220153808594]], 'keypoints': [[145.0775909423828, 55.070377349853516, 167.0526123046875, 51.27381134033203, 155.4307403564453, 67.80669403076172, 152.44337463378906, 78.16500091552734, 170.25086975097656, 75.12496948242188]]}
          boxes = result['boxes']
          new_image = Image.new("RGBA", image_item.size, (255, 255, 255, 0))
          draw = ImageDraw.Draw(new_image)
          for box in boxes:
              # box 格式为 [x, y, x1, y1]
              x, y, x1, y1 = box
              # 使用 ImageDraw 对象绘制矩形
              draw.rectangle([x, y, x1, y1], fill=(0, 0, 0, 255))

          out_tensor = pil2tensor(new_image)
          res_images.extend(out_tensor)
        return (res_images,)
    
